#!/usr/bin/env python
# Copyright 2023 Statistics and Machine Learning Research Group at HKUST. All rights reserved.
"""A one-line summary of the module or program, terminated by a period.

Leave one blank line.  The rest of this docstring should contain an
overall description of the module or program.  Optionally, it may also
contain a brief description of exported classes and functions and/or usage
examples.

Typical usage example:

  foo = ClassFoo()
  bar = foo.FunctionBar()
"""

import os
import sys

sys.path.remove(os.path.abspath(os.path.dirname(sys.argv[0])))
from transformers import HfArgumentParser

from lmflow.args import (
    AutoArguments,
    DatasetArguments,
    ModelArguments,
)
from lmflow.datasets.dataset import Dataset
from lmflow.models.auto_model import AutoModel
from lmflow.pipeline.auto_pipeline import AutoPipeline


def main():
    # Parses arguments
    pipeline_name = "finetuner"
    # 根据pipeline_name返回一个保存参数的类的实例；这里是返回一个微调类的实例，该类中保存了很多微调模型需要的参数
    PipelineArguments = AutoArguments.get_pipeline_args_class(pipeline_name)
    
    # 从命令行中读取参数，自动识别输入的参数属于哪个类
    parser = HfArgumentParser((ModelArguments, DatasetArguments, PipelineArguments))
    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
        # If we pass only one argument to the script and it's the path to a json file,
        # let's parse it to get our arguments.
        model_args, data_args, pipeline_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
    else:
        # 跳转到这里，model_args, data_args, pipeline_args分别保存了一些参数
        model_args, data_args, pipeline_args = parser.parse_args_into_dataclasses()

    # Initialization
    # 根据pipeline_name返回某个任务类的实例，这里是返回一个微调任务类的实例
    finetuner = AutoPipeline.get_pipeline(
        pipeline_name=pipeline_name,
        model_args=model_args,
        data_args=data_args,
        pipeline_args=pipeline_args,
    )
    dataset = Dataset(data_args)
    model = AutoModel.get_model(model_args)

    # Finetuning
    # 开始微调
    finetuner.tune(model=model, dataset=dataset)


if __name__ == "__main__":
    main()
